import torch.nn as nn
import torch
import torch.nn.functional as F


class SiameseNetwork(nn.Module):
    def __init__(self, vocab_size, d_model):
        super(SiameseNetwork, self).__init__()
        self.vocab_size = vocab_size
        self.d_model = d_model

        self.embedding = nn.Embedding(vocab_size, d_model)
        self.LSTM = nn.LSTM(d_model, d_model)
        self.mean = lambda x: torch.mean(x, dim=1)

    def Mynorm(self, x):
        return x / torch.sqrt(torch.sum(x**2, axis=-1, keepdim=True))

    def single_forward(self, x):
        x = self.embedding(x)
        x, _ = self.LSTM(x)
        x = self.mean(x)
        x = self.Mynorm(x)
        return x

    def forward(self, x1, x2):
        x1 = self.single_forward(x1)
        x2 = self.single_forward(x2)
        return x1, x2


class TripletLoss(torch.nn.Module):
    def __init__(self, device, margin=0.25):
        super(TripletLoss, self).__init__()
        self.margin = margin
        self.device = device

    def forward(self, x1, x2):
        similarity = torch.matmul(x1, x2.T)
        batch_size = similarity.shape[0]
        eye = torch.eye(batch_size).to(self.device)
        cost_ap = torch.diagonal(similarity)
        mean_neg = ((1 - eye) * similarity).sum(axis=1) / (batch_size - 1)
        closest_neg = (torch.matmul(x1, x2.T) - 2 * eye).max(axis=1).values
        loss1 = F.relu(self.margin - cost_ap + mean_neg)
        loss2 = F.relu(self.margin - cost_ap + closest_neg)
        triplet_loss = torch.mean(loss1 + loss2)
        return triplet_loss
